import os
import json
import argparse
import serial
from typing import Dict, Any, List
import DM_CAN as dm


# 默认臂配置（参考 motor_config.py），保存到 arm_config.json
DEFAULT_ARM_CONFIG: Dict[str, Any] = {
    'motors': [
        {'type': 'DM10010L', 'slave_id': 0x01, 'master_id': 0x11},
        {'type': 'DM10010L', 'slave_id': 0x02, 'master_id': 0x12},
        {'type': 'DM4340_48V', 'slave_id': 0x03, 'master_id': 0x13},
        {'type': 'DM4340_48V', 'slave_id': 0x04, 'master_id': 0x14},
        {'type': 'DM4310_48V', 'slave_id': 0x05, 'master_id': 0x15},
        {'type': 'DM4310_48V', 'slave_id': 0x06, 'master_id': 0x16},
        {'type': 'DM4310_48V', 'slave_id': 0x07, 'master_id': 0x17},
    ],
    'serial': {
        'port': 'COM6',
        'baudrate': 921600,
        'timeout': 0.5
    },
    'display': {
        'refresh_rate': 30,
        'position_warning_threshold': 90,
        'position_danger_threshold': 150,
        'velocity_warning_threshold': 15,
        'velocity_danger_threshold': 30,
        'torque_warning_threshold': 4,
        'torque_danger_threshold': 7
    },
    # 当前默认臂侧（记录未指定 arm 时使用）
    'arm_side': 'left',
    # 记录当前各电机位置（由 record_motor_pos 写入），按左右臂分桶
    'motor_positions': {
        'left': {
            # '0x01': 0.0
        },
        'right': {
        }
    },
    # 记录每个电机的限位（min/max），按左右臂分桶
    'motor_limits': {
        'left': {
            # '0x01': {'min': 0.0, 'max': 1.57}
        },
        'right': {
        }
    }
}

MOTOR_TYPE_MAP = {
    'DM10010L': dm.DM_Motor_Type.DM10010L,
    'DM4340_48V': dm.DM_Motor_Type.DM4340_48V,
    'DM4310_48V': dm.DM_Motor_Type.DM4310_48V,
}


def _default_config_path() -> str:
    # 将配置文件固定在当前脚本目录下
    return os.path.join(os.path.dirname(__file__), 'arm_config.json')


def _load_json_config(path: str) -> Dict[str, Any]:
    if os.path.exists(path):
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)
    # 首次运行：生成默认配置
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(DEFAULT_ARM_CONFIG, f, ensure_ascii=False, indent=2)
    return json.loads(json.dumps(DEFAULT_ARM_CONFIG))


def _save_json_config(path: str, cfg: Dict[str, Any]) -> None:
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)


def _get_by_key_path(cfg: Dict[str, Any], key_path: str):
    keys = key_path.split('.')
    cur = cfg
    for k in keys:
        if isinstance(cur, list) and k.isdigit():
            cur = cur[int(k)]
        else:
            cur = cur[k]
    return cur


def _set_by_key_path(cfg: Dict[str, Any], key_path: str, value: Any):
    keys = key_path.split('.')
    cur = cfg
    for k in keys[:-1]:
        if isinstance(cur, list) and k.isdigit():
            cur = cur[int(k)]
        else:
            cur = cur[k]
    last = keys[-1]
    if isinstance(cur, list) and last.isdigit():
        cur[int(last)] = value
    else:
        cur[last] = value


def _ensure_side_buckets(cfg: Dict[str, Any]) -> None:
    """确保配置中的 motor_positions 和 motor_limits 具备左右臂分桶结构。

    - 若检测到旧版平铺结构（如 {'0x01': 0.0}），则默认归入 left 桶。
    - 若缺失键，则补齐空结构。
    """
    cfg.setdefault('arm_side', 'left')

    mp = cfg.get('motor_positions')
    if not isinstance(mp, dict) or ('left' not in mp or 'right' not in mp):
        if isinstance(mp, dict) and any(isinstance(k, str) and k.startswith('0x') for k in mp.keys()):
            cfg['motor_positions'] = {'left': mp, 'right': {}}
        else:
            cfg['motor_positions'] = {'left': {}, 'right': {}}

    ml = cfg.get('motor_limits')
    if not isinstance(ml, dict) or ('left' not in ml or 'right' not in ml):
        if isinstance(ml, dict) and any(isinstance(k, str) and k.startswith('0x') for k in ml.keys()):
            cfg['motor_limits'] = {'left': ml, 'right': {}}
        else:
            cfg['motor_limits'] = {'left': {}, 'right': {}}


class Armcalibration:
    """
    机械臂校准类：
    - set_single_motor_zero: 调用库函数设置单个电机零位
    - set_arm_zero: 预留（等待后续补充）
    - record_motor_pos: 读取电机位置并写入 arm_config.json
    - set_arm_config: 设置 JSON 配置参数
    - set_arm_x/set_arm_y/set_arm_z: 预留（等待后续补充）
    - set_single_motor_pos: 参考 controlMIT 设置单个电机位置
    """

    def __init__(self, config_file: str = None):
        self.config_file = config_file or _default_config_path()
        self.config: Dict[str, Any] = _load_json_config(self.config_file)
        _ensure_side_buckets(self.config)

        # 初始化串口与控制器
        serial_cfg = self.config['serial']
        self.serial_device = serial.Serial(
            serial_cfg['port'], serial_cfg['baudrate'], timeout=serial_cfg['timeout']
        )
        self.controller = dm.MotorControl(self.serial_device)

        # 初始化电机对象并建立映射
        self.motors: List[dm.Motor] = []
        self.motors_by_id: Dict[int, dm.Motor] = {}
        for m in self.config['motors']:
            mtype = MOTOR_TYPE_MAP[m['type']]
            motor = dm.Motor(mtype, m['slave_id'], m['master_id'])
            self.motors.append(motor)
            self.motors_by_id[motor.SlaveID] = motor
            self.controller.addMotor(motor)

    def close(self):
        # 失能电机并关闭串口
        try:
            for m in self.motors:
                self.controller.disable(m)
        finally:
            if self.serial_device:
                try:
                    self.serial_device.close()
                except Exception:
                    pass

    # ---- 功能函数 ----
    def set_single_motor_zero(self, slave_id: int) -> None:
        """设置单个电机零位（建议在电机失能状态下执行）。"""
        motor = self.motors_by_id.get(slave_id)
        if motor is None:
            print(f"Motor {hex(slave_id)} not found")
            return
        # 按照 README 建议，先失能再设置零位
        self.controller.disable(motor)
        self.controller.set_zero_position(motor)
        print(f"Set zero for motor {hex(slave_id)} done.")

    def set_arm_zero(self):
        """整臂零位设置（预留）。"""
        print("set_arm_zero: TODO - 等待后续补充")

    def _is_motor_responsive(self, motor: dm.Motor) -> bool:
        """检查电机是否有响应：通过读取电机参数判断通讯是否正常。"""
        try:
            # 优先读取软件版本；若无响应，尝试读取硬件版本
            ver = self.controller.read_motor_param(motor, dm.DM_variable.sw_ver)
            if ver is None:
                ver2 = self.controller.read_motor_param(motor, dm.DM_variable.hw_ver)
                return ver2 is not None
            return True
        except Exception:
            return False

    def record_motor_pos(self, slave_id: int, limit: str = None, direction: str = None, arm: str = None) -> None:
        """读取电机位置并记录到 arm_config.json。

        Args:
            slave_id: 电机 SlaveID
            limit: 可选，'min' 或 'max'；用于标记该位置为该电机的最小/最大限位。
            direction: 可选，'cw' 或 'ccw'；当检测到两个限位位置相同或近似（<0.1）时，必须提供旋转方向。
            arm: 可选，'left' 或 'right'；不提供时使用配置中的 'arm_side'。
        """
        motor = self.motors_by_id.get(slave_id)
        if motor is None:
            print(f"Motor {hex(slave_id)} not found")
            return
        # 读数前先使能电机，刷新状态，读取位置；最后失能
        try:
            self.controller.enable(motor)
            # 刷新状态，并检查电机是否有响应；无响应则不记录
            self.controller.refresh_motor_status(motor)
            if not self._is_motor_responsive(motor):
                print(f"Motor {hex(slave_id)} 未响应，无法记录。请检查供电、串口、CAN 适配器、ID 配置。")
                return
            pos = motor.getPosition()
        finally:
            try:
                self.controller.disable(motor)
            except Exception:
                pass
        # 写入 JSON（按左右臂分桶）
        _ensure_side_buckets(self.config)
        arm_side = arm or self.config.get('arm_side', 'left')
        if arm_side not in ('left', 'right'):
            print(f"无效的臂侧: {arm_side}，仅支持 left 或 right")
            return
        pos_value = float(pos)
        sid = hex(slave_id)
        self.config['motor_positions'].setdefault(arm_side, {})
        self.config['motor_positions'][arm_side][sid] = pos_value

        # 如果提供了限位标志，则同时记录到 motor_limits
        if limit in ('min', 'max'):
            self.config.setdefault('motor_limits', {})
            self.config['motor_limits'].setdefault(arm_side, {})
            self.config['motor_limits'][arm_side].setdefault(sid, {})
            self.config['motor_limits'][arm_side][sid][limit] = pos_value

            # 检测是否为360度电机（两个标记位置相同或近似）
            mlims = self.config['motor_limits'][arm_side][sid]
            if 'min' in mlims and 'max' in mlims:
                if abs(mlims['max'] - mlims['min']) < 0.1:
                    # 标记为360度电机
                    mlims['is_360'] = True
                    # 方向校验与记录：在该情况下必须提供方向
                    if direction is None:
                        _save_json_config(self.config_file, self.config)
                        raise ValueError(
                            f"Motor {sid} 判定为360度电机，请使用 --dir 指定旋转方向（cw/ccw）"
                        )
                    if direction not in ('cw', 'ccw'):
                        _save_json_config(self.config_file, self.config)
                        raise ValueError(
                            f"无效的旋转方向: {direction}，仅支持 'cw' 或 'ccw'"
                        )
                    mlims['rotation_direction'] = direction
                else:
                    # 非360度，清理可能存在的标记
                    mlims['is_360'] = False
                    mlims.pop('rotation_direction', None)
        _save_json_config(self.config_file, self.config)
        if limit in ('min', 'max'):
            base_msg = f"Recorded motor {sid} ({arm_side}) position: {pos_value} as {limit} limit"
            if self.config['motor_limits'].get(sid, {}).get('is_360'):
                dir_str = self.config['motor_limits'][sid].get('rotation_direction')
                if dir_str:
                    print(base_msg + f" (360°, dir={dir_str})")
                else:
                    print(base_msg + " (360°, 未设置旋转方向)")
            else:
                print(base_msg)
        else:
            print(f"Recorded motor {sid} ({arm_side}) position: {pos_value}")

    def set_arm_config(self, key_path: str, value: Any, save: bool = True) -> None:
        """设置 arm_config.json 配置参数（参考 motor_config.py 的键路径风格）。"""
        _set_by_key_path(self.config, key_path, value)
        if save:
            _save_json_config(self.config_file, self.config)
        print(f"Config set: {key_path} = {value}")

    def set_arm_x(self, value: Any):
        print("set_arm_x: TODO - 等待后续补充")

    def set_arm_y(self, value: Any):
        print("set_arm_y: TODO - 等待后续补充")

    def set_arm_z(self, value: Any):
        print("set_arm_z: TODO - 等待后续补充")

    def set_single_motor_pos(self, slave_id: int, kp: float, kd: float, q: float, dq: float, tau: float) -> None:
        """参考库中 controlMIT，设置单个电机位置/速度/力矩（MIT模式）。"""
        motor = self.motors_by_id.get(slave_id)
        if motor is None:
            print(f"Motor {hex(slave_id)} not found")
            return
        # 可选：确保控制模式为 MIT；有的固件不需要，保留简化实现
        # self.controller.switchControlMode(motor, dm.Control_Type.MIT)
        # 使能电机后下发控制
        self.controller.enable(motor)
        self.controller.controlMIT(motor, kp, kd, q, dq, tau)
        print(
            f"Sent MIT command to motor {hex(slave_id)}: kp={kp}, kd={kd}, q={q}, dq={dq}, tau={tau}"
        )


def _build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description='T9 Arm Calibration CLI')
    parser.add_argument('--interactive', '-i', action='store_true', help='交互模式')
    sub = parser.add_subparsers(dest='cmd')

    sub.add_parser('list', help='列出配置中的电机')

    p_zero = sub.add_parser('zero', help='设置单个电机零位')
    p_zero.add_argument('--id', type=lambda x: int(x, 0), required=True, help='电机 SlaveID，如 0x01 或 1')

    p_record = sub.add_parser('record', help='记录单个电机当前位置到 JSON，可带限位标记与旋转方向')
    p_record.add_argument('--id', type=lambda x: int(x, 0), required=True, help='电机 SlaveID，如 0x01 或 1')
    grp = p_record.add_mutually_exclusive_group()
    grp.add_argument('--min', action='store_true', help='将该位置标记为该电机最小限位')
    grp.add_argument('--max', action='store_true', help='将该位置标记为该电机最大限位')
    p_record.add_argument('--dir', choices=['cw', 'ccw'], help='旋转方向：cw=顺时针，ccw=逆时针；仅在两限位位置相近时必填')
    p_record.add_argument('--arm', choices=['left', 'right'], help='臂侧：left=左臂，right=右臂；不提供时使用配置中的 arm_side')

    p_cfg = sub.add_parser('set-config', help='设置臂配置参数（key.path = value）')
    p_cfg.add_argument('--key', required=True, help='键路径，如 serial.port 或 motors.0.type')
    p_cfg.add_argument('--value', required=True, help='值（会尝试按 int/float/str 解析）')

    p_pos = sub.add_parser('set-pos', help='MIT 控制设置单个电机位置')
    p_pos.add_argument('--id', type=lambda x: int(x, 0), required=True, help='电机 SlaveID，如 0x01 或 1')
    p_pos.add_argument('--kp', type=float, required=True)
    p_pos.add_argument('--kd', type=float, required=True)
    p_pos.add_argument('--q', type=float, required=True)
    p_pos.add_argument('--dq', type=float, required=True)
    p_pos.add_argument('--tau', type=float, required=True)

    sub.add_parser('arm-zero', help='整臂零位设置（预留）')
    p_ax = sub.add_parser('set-x', help='设置 Arm X（预留）')
    p_ax.add_argument('--value', required=True)
    p_ay = sub.add_parser('set-y', help='设置 Arm Y（预留）')
    p_ay.add_argument('--value', required=True)
    p_az = sub.add_parser('set-z', help='设置 Arm Z（预留）')
    p_az.add_argument('--value', required=True)

    return parser


def _auto_parse_value(s: str):
    # 将字符串转换为更合适的类型
    for cast in (lambda v: int(v, 0), float, str):
        try:
            return cast(s)
        except Exception:
            continue
    return s


def _prompt_int(msg: str) -> int:
    while True:
        s = input(msg).strip()
        try:
            return int(s, 0)
        except Exception:
            print("输入无效，请输入整数（支持十六进制如 0x01）")


def _prompt_float(msg: str) -> float:
    while True:
        s = input(msg).strip()
        try:
            return float(s)
        except Exception:
            print("输入无效，请输入浮点数")


def _prompt_choice(msg: str, choices, allow_empty=False, default=None):
    choices_lower = [c.lower() for c in choices]
    while True:
        s = input(msg).strip().lower()
        if not s:
            if allow_empty:
                return default
            print("输入不能为空")
            continue
        if s in choices_lower:
            return s
        print(f"无效选择，请输入 {choices}")


def _interactive_loop(ac: Armcalibration):
    print("进入交互模式。按编号选择操作，输入 0 退出。")
    while True:
        print("\n=== 菜单 ===")
        print("1) 列出电机")
        print("2) 设置单电机零位")
        print("3) 记录电机位置（支持限位与旋转方向）")
        print("4) 修改配置参数")
        print("5) MIT 控制设置单电机位置")
        print("6) 整臂零位（预留）")
        print("7) 设置 Arm X（预留）")
        print("8) 设置 Arm Y（预留）")
        print("9) 设置 Arm Z（预留）")
        print("0) 退出")

        choice = input("请选择操作编号: ").strip()

        if choice == '0':
            print("退出交互模式。")
            break
        elif choice == '1':
            print('已配置电机:')
            for m in ac.motors:
                print(f"- type={m.MotorType.name}, slave_id={hex(m.SlaveID)}, master_id={hex(m.MasterID)}")
        elif choice == '2':
            sid = _prompt_int("输入电机 SlaveID（如 0x01 或 1）: ")
            ac.set_single_motor_zero(sid)
        elif choice == '3':
            sid = _prompt_int("输入电机 SlaveID（如 0x01 或 1）: ")
            lim = _prompt_choice("是否设置限位（none/min/max）: ", ['none', 'min', 'max'])
            direction = None
            if lim in ('min', 'max'):
                dir_opt = _prompt_choice("旋转方向（cw/ccw，留空跳过）：", ['cw', 'ccw'], allow_empty=True, default=None)
                direction = dir_opt if dir_opt else None
            arm_opt = _prompt_choice("臂侧（left/right，留空默认当前配置）：", ['left', 'right'], allow_empty=True, default=None)
            arm_side = arm_opt if arm_opt else None
            try:
                ac.record_motor_pos(sid, None if lim == 'none' else lim, direction, arm_side)
            except ValueError as e:
                print(f"错误：{e}")
                # 若因 360° 电机强制需要方向，则再次提示方向
                if '旋转方向' in str(e):
                    dir_force = _prompt_choice("请指定旋转方向（cw/ccw）：", ['cw', 'ccw'])
                    ac.record_motor_pos(sid, None if lim == 'none' else lim, dir_force, arm_side)
        elif choice == '4':
            key = input("输入键路径（如 serial.port 或 motors.0.type）：").strip()
            val_str = input("输入值（自动按 int/float/str 解析）：").strip()
            val = _auto_parse_value(val_str)
            ac.set_arm_config(key, val, save=True)
        elif choice == '5':
            sid = _prompt_int("输入电机 SlaveID（如 0x01 或 1）: ")
            kp = _prompt_float("输入 kp: ")
            kd = _prompt_float("输入 kd: ")
            q = _prompt_float("输入 q: ")
            dq = _prompt_float("输入 dq: ")
            tau = _prompt_float("输入 tau: ")
            ac.set_single_motor_pos(sid, kp, kd, q, dq, tau)
        elif choice == '6':
            ac.set_arm_zero()
        elif choice == '7':
            val = input("输入 Arm X 值：").strip()
            ac.set_arm_x(val)
        elif choice == '8':
            val = input("输入 Arm Y 值：").strip()
            ac.set_arm_y(val)
        elif choice == '9':
            val = input("输入 Arm Z 值：").strip()
            ac.set_arm_z(val)
        else:
            print("无效编号，请重新选择。")


def main():
    parser = _build_arg_parser()
    args = parser.parse_args()
    ac = Armcalibration()

    try:
        # 交互模式优先，或未提供子命令时进入交互模式
        if getattr(args, 'interactive', False) or not getattr(args, 'cmd', None):
            _interactive_loop(ac)
        elif args.cmd == 'list':
            print('Configured Motors:')
            for m in ac.motors:
                print(f"- type={m.MotorType.name}, slave_id={hex(m.SlaveID)}, master_id={hex(m.MasterID)}")

        elif args.cmd == 'zero':
            ac.set_single_motor_zero(args.id)

        elif args.cmd == 'record':
            limit_flag = 'min' if getattr(args, 'min', False) else ('max' if getattr(args, 'max', False) else None)
            ac.record_motor_pos(args.id, limit_flag, getattr(args, 'dir', None), getattr(args, 'arm', None))

        elif args.cmd == 'set-config':
            val = _auto_parse_value(args.value)
            ac.set_arm_config(args.key, val, save=True)

        elif args.cmd == 'set-pos':
            ac.set_single_motor_pos(args.id, args.kp, args.kd, args.q, args.dq, args.tau)

        elif args.cmd == 'arm-zero':
            ac.set_arm_zero()

        elif args.cmd == 'set-x':
            ac.set_arm_x(args.value)

        elif args.cmd == 'set-y':
            ac.set_arm_y(args.value)

        elif args.cmd == 'set-z':
            ac.set_arm_z(args.value)

        else:
            parser.print_help()
    finally:
        ac.close()


if __name__ == '__main__':
    main()